#include "memory.h"
#include "print.h"

#define PAGE_SIZE 4096

#define MEM_BITMAP_BASE 0xc009a000

#define PDE_IDX(addr) (((uint32_t)addr & 0xffc00000) >> 22)
#define PTE_IDX(addr) (((uint32_t)addr & 0x003ff000) >> 12)

#define PAGE_MEM (256 * PAGE_SIZE)
#define HEAP_START (0xc0100000 + PAGE_MEM)
#define VIRTUAL_START 0xc0100000

struct mem_pool kernel_mem_pool;
struct mem_pool user_mem_pool;
struct virtual_pool kernel_virtual_pool;

void memory_init(void)
{
	uint32_t total_mem_size = *(uint32_t*)0xc0000910;
	
	uint32_t kernel_mem_size = total_mem_size / 2 / PAGE_SIZE * PAGE_SIZE;
	uint32_t user_mem_size = total_mem_size / 2 / PAGE_SIZE * PAGE_SIZE;
	
	uint32_t kernel_bitmap_size = kernel_mem_size / PAGE_SIZE / 8;
	uint32_t user_bitmap_size = user_mem_size / PAGE_SIZE / 8;
	uint32_t virtual_bitmap_size = kernel_bitmap_size;
	
	void* kernel_bitmap_addr = (void*)MEM_BITMAP_BASE;
	void* user_bitmap_addr = kernel_bitmap_addr + kernel_bitmap_size;
	void* virtual_bitmap_addr = user_bitmap_addr + user_bitmap_size;
	
	void* kernel_paddr_start = (void*)HEAP_START;
	void* user_paddr_start = kernel_paddr_start + kernel_mem_size;
	void* virtual_paddr_start = (void*)VIRTUAL_START;
	
	print_string("total mem: ");
	print_hex(total_mem_size);
	print_char('\n');
	
	bitmap_init(&kernel_mem_pool.addr_pool, (uint8_t*)kernel_bitmap_addr, kernel_bitmap_size);
	kernel_mem_pool.paddr_start = kernel_paddr_start;
	kernel_mem_pool.size = kernel_mem_size;
	
	bitmap_init(&user_mem_pool.addr_pool, (uint8_t*)user_bitmap_addr, user_bitmap_size);
	user_mem_pool.paddr_start = user_paddr_start;
	user_mem_pool.size = user_mem_size;
	
	bitmap_init(&kernel_virtual_pool.addr_pool, (uint8_t*)virtual_bitmap_addr, virtual_bitmap_size);
	kernel_virtual_pool.vaddr_start = virtual_paddr_start;
}

static void* malloc_virtual_page(enum pool_flags pf, uint32_t pg_count)
{
	int32_t index = -1;
	int i;
	void* ret = NULL;
	
	if (pg_count == 0)
		return NULL;
	
	if (pf == PF_KERNEL)
	{
		index = bitmap_scan_test(&kernel_virtual_pool.addr_pool, pg_count);
		
		if (index >= 0)
		{
			for (i = 0; i < pg_count; ++i)
			{
				bitmap_set(&kernel_virtual_pool.addr_pool, index + i);
			}
			
			ret = kernel_virtual_pool.vaddr_start + PAGE_SIZE * index;
		}
	}
	else if (pf == PF_USER)
	{
	
	}
	
	return ret;
}

static int8_t add_page(void* virtual, void* physical)
{
	int32_t index;
	uint32_t* pde = (uint32_t*)(0xFFFFF000 | PDE_IDX(virtual) * 4);
	uint32_t* pte = (uint32_t*)(0xFFC00000 | PDE_IDX(virtual) << 12 | PTE_IDX(virtual) * 4);
	
	if (NULL != physical)
	{
		if (!(*pde & 0x1))
		{
			index = bitmap_scan_test(&kernel_mem_pool.addr_pool, 1);
			if (index < 0)
				return -1;
			
			bitmap_set(&kernel_mem_pool.addr_pool, index);
			
			void* paddr = kernel_mem_pool.paddr_start + PAGE_SIZE * index;
			*pde = (uint32_t)paddr & 0xFFFF000 | PG_P_1 | PG_RW_W | PG_US_U;
		}
		
		*pte = (uint32_t)physical & 0xFFFF000 | PG_P_1 | PG_RW_W | PG_US_U;
	}
	else
	{
		*pte = 0;
	}
	
	return 0;
}

void* malloc_page(enum pool_flags pf, uint32_t pg_count)
{
	int i;
	int32_t index;
	
	if (pg_count == 0)
		return NULL;
	
	void* virtual_addr = malloc_virtual_page(pf, pg_count);
	if (NULL == virtual_addr)
	{
		return NULL;
	}
	
	struct mem_pool* pool = pf == PF_KERNEL ? &kernel_mem_pool : &user_mem_pool;
	
	for (i = 0; i < pg_count; ++i)
	{
		index = bitmap_scan_test(&pool->addr_pool, 1);
		if (index >= 0)
		{
			bitmap_set(&pool->addr_pool, index);
			
			void* physical_addr = pool->paddr_start + PAGE_SIZE * index;
			if (add_page(virtual_addr, physical_addr) < 0)
			{
				// TODO rollback
				return NULL;
			}
		}
		else
		{
			// TODO rollback
			return NULL;
		}
	}
	
	return virtual_addr;
}
